## 从标签文件中筛选出所有 <name> 只要有指定类别的标签文件及对应图片文件，并将它们复制到新的目标文件夹中
import os
import shutil
import xml.etree.ElementTree as ET

def filter_files_by_label(src_folder, dest_folder, allowed_categories):
    # 定义源文件夹和目标文件夹
    annotations_src = os.path.join(src_folder, "Annotations")
    images_src = os.path.join(src_folder, "JPEGImages")

    annotations_dest = os.path.join(dest_folder, "Annotations")
    images_dest = os.path.join(dest_folder, "JPEGImages")

    # 创建目标文件夹和子文件夹
    os.makedirs(annotations_dest, exist_ok=True)
    os.makedirs(images_dest, exist_ok=True)

    # 遍历标签文件
    for filename in os.listdir(annotations_src):
        if not filename.endswith('.xml'):
            continue

        xml_path = os.path.join(annotations_src, filename)

        # 解析XML文件
        tree = ET.parse(xml_path)
        root = tree.getroot()

        # 检查<name>标签是否都在允许的类别中
        valid = False
        for obj in root.findall('object'):
            name = obj.find('name').text
            if name in allowed_categories:
                valid = True
                break

        # 如果标签文件符合要求，复制标签文件和对应的图片文件
        if valid:
            # 复制标签文件
            shutil.copy(xml_path, os.path.join(annotations_dest, filename))

            # 找到对应的图片文件
            image_filename = filename.replace('.xml', '.jpg')
            image_path = os.path.join(images_src, image_filename)

            if os.path.exists(image_path):
                shutil.copy(image_path, os.path.join(images_dest, image_filename))

# 示例使用
src_folder = "/home/shuai/Documents/mydataset/my_usv_dock_exp/VOCdevkit_USV"  # 替换为你的源文件夹路径
dest_folder = "/home/shuai/Documents/mydataset/my_usv_dock_exp/VOCdevkit_USV_exp"  # 替换为你的目标文件夹路径
allowed_categories = ["USV"]  # 替换为你指定的类别

filter_files_by_label(src_folder, dest_folder, allowed_categories)
